package com.consume.controller;

import brave.propagation.CurrentTraceContext;
import brave.propagation.TraceContext;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.concurrent.Callable;

/**
 * log日志跨线程共用traceId；获取父线程的traceId失败,不打印traceId
 * 
 * @author zhangbowen@heytea.com
 * @since 2019-08-05 09:14
 */
public class TraceCrossWarp implements InvocationHandler {
    private Object target;
    private TraceContext traceContext;

    public static Runnable of(Runnable runnable) {
        TraceCrossWarp warp = new TraceCrossWarp(runnable);
        return warp.getRunnableProxy();
    }

    public static <T> Callable<T> of(Callable<T> callable) {
        TraceCrossWarp warp = new TraceCrossWarp(callable);
        return warp.getCallableProxy();
    }

    private TraceCrossWarp(Object target) {
        this.target = target;
        this.traceContext = CurrentTraceContext.Default.create().get();
    }

    private Runnable getRunnableProxy() {
        return (Runnable)Proxy.newProxyInstance(this.getClass().getClassLoader(), target.getClass().getInterfaces(),
            this);
    }

    private <T> Callable<T> getCallableProxy() {
        return (Callable<T>)Proxy.newProxyInstance(this.getClass().getClassLoader(), target.getClass().getInterfaces(),
            this);
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        try (CurrentTraceContext.Scope scope = TraceScope.newThreadCrossScope(traceContext)) {
            return method.invoke(target, args);
        }
    }
}
